from typing import Dict, List, Type
from .models import (
    ConcreteValueType,
    Variable,
    ArrayValue,
    StructValue,
    PointerValue,
    FunctionValue,
    ArrayInitializerValue,
    _shallow_copy_variable,
)
from ..universal_ast_nodes import *
from ..universal_cfg_extractor import CFGBuilder
from ..uast_queries import UASTQuery


class FunctionReturn(StopIteration):
    def __init__(self, ret_value: List[ConcreteValueType]) -> None:
        self.ret_value = ret_value
        super().__init__()


# def create_struct_value_by_decl(struct_decl: StructDecl) -> StructValue:
#     return StructValue(
#         {
#             member.name: create_concreate_value_by_type(member.type)
#             for member in struct_decl.members
#         }
#     )


def create_concreate_value_by_type(type: DATA_TYPE) -> ConcreteValueType:
    match type:
        case IntType():
            return 0
        case FloatType():
            return 0.0
        case StructType():
            raise NotImplementedError
        case _:
            raise NotImplementedError(f"Type {type} is not implemented")


class BaseUASTExecutor:
    def __init__(
        self,
        variables: List[Dict[str, Variable]],
        method_asts: Dict[str, MethodDecl] = {},
    ) -> None:
        self.variables: List[Dict[str, Variable]] = variables
        self.method_asts: Dict[str, MethodDecl] = method_asts
        self.evaluator_mapping: Dict[
            Type[SourceElement], Callable[[SourceElement], Optional[ConcreteValueType]]
        ] = {
            CompoundDecl: self._eval_compound_decl,
            MethodDecl: self._eval_method_decl,
            VarDecl: self._eval_var_decl,
            BinaryExpr: self._eval_binary_expr,
            Name: self._eval_name,
            UnaryExpr: self._eval_unary_expr,
            FieldAccessExpr: self._eval_field_access_expr,
            ArrayAccessExpr: self._eval_array_access_expr,
            CallExpr: self._eval_call_expr,
            ReturnStmt: self._eval_return_stmt,
            Assignment: self._eval_assignment,
            Literal: self._eval_literal,
            ArrayInitializer: self._eval_array_initializer,
        }

    def _is_concrete_value(self, value: Any) -> bool:
        if isinstance(value, (int, float, bool, str)) or value is None:
            return True
        else:
            return False

    def _assign_value_to_target(self, l_value: Variable, r_value: Any):
        if isinstance(l_value.value, ArrayValue):
            for i, v in enumerate(l_value.value.values):
                self._assign_value_to_target(v, r_value[i].value)
        elif isinstance(l_value.value, StructValue):
            for i, item in enumerate(l_value.value.attributes.items()):
                k, v = item
                self._assign_value_to_target(v, r_value[i].value)
        elif isinstance(l_value.value, PointerValue):
            assert isinstance(r_value, PointerValue), r_value
            l_value.value.pointee = r_value.pointee
        elif self._is_concrete_value(l_value.value):
            l_value.value = r_value
        else:
            raise NotImplementedError(l_value.value)

    def assign_value(self, l_value: Variable, r_value: Any):
        if isinstance(r_value, (StructValue, ArrayValue)):
            r_value = r_value.shallow_copy()
        self._assign_value_to_target(l_value, r_value)

    def get_method_uast(self, node: SourceElement) -> MethodDecl:
        match node:
            case Name():
                variable = self._eval_name(node)
                assert isinstance(variable, FunctionValue), variable
                return self.method_asts[variable.name]
            case _:
                raise NotImplementedError(node)

    def eval_l_value(self, node: SourceElement) -> Variable:
        match node:
            case Name(id=id):
                for level in range(len(self.variables) - 1, -1, -1):
                    if id in self.variables[level]:
                        return self.variables[level][id]
            case UnaryExpr(sign="*", op_before_expr=True):
                l_value = self.eval_l_value(node.expression)
                assert isinstance(l_value.value, PointerValue), l_value
                return l_value.value.pointee
            case ArrayAccessExpr(index=index, target=target):
                target_arr: ArrayValue = self.eval_l_value(target).value
                index_value = self._eval_node(index)
                return target_arr.values[index_value]
            case FieldAccessExpr(name=name, target=target):
                l_value = self.eval_l_value(target).value
                if isinstance(l_value, PointerValue):
                    target_object: StructValue = l_value.pointee.value
                else:
                    target_object: StructValue = l_value
                assert isinstance(target_object, StructValue), target_object
                return target_object.attributes[name]
            case _:
                raise NotImplementedError(node)

    def _eval_literal(self, node: Literal) -> ConcreteValueType:
        match node.kind:
            case "int":
                if node.value.lower().startswith("0x"):
                    return int(node.value[2:], 16)
                elif node.value.lower().startswith("0b"):
                    return int(node.value[2:], 2)
                return int(node.value)
            case _:
                raise NotImplementedError(node.kind)

    def _eval_array_initializer(self, node: ArrayInitializer) -> ConcreteValueType:
        values = []
        for element in node.elements:
            values.append(self._eval_node(element))
        return ArrayInitializerValue(values)

    def _eval_name(self, node: Name) -> ConcreteValueType:
        for level in range(len(self.variables) - 1, -1, -1):
            if node.id in self.variables[level]:
                return self.variables[level][node.id].value
        print(self.variables)
        raise ValueError(f"Variable {node.id} not found")

    def _eval_field_access_expr(self, node: FieldAccessExpr) -> ConcreteValueType:
        target = self._eval_node(node.target)
        name = node.name
        if isinstance(target, PointerValue):
            target_object: StructValue = target.pointee.value
        else:
            target_object: StructValue = target
        assert isinstance(target_object, StructValue), target_object
        return target_object.attributes[name].value

    def _eval_array_access_expr(self, node: ArrayAccessExpr) -> ConcreteValueType:
        target = self._eval_node(node.target)
        index = self._eval_node(node.index)
        target_object: ArrayValue = target
        assert isinstance(target_object, ArrayValue), target_object
        return target_object.values[index].value

    def _eval_call_expr(self, node: CallExpr) -> ConcreteValueType:
        # match node:
        method_uast = self.get_method_uast(node.name)
        stack_variables = {}
        param_names = (
            UASTQuery.get_all_params(method_uast).map(lambda node: node.name.id).l
        )
        arg_values = [self._eval_node(item) for item in node.arguments]
        for name, value in zip(param_names, arg_values):
            stack_variables[name] = _shallow_copy_variable(Variable(value))
        self.variables.append(stack_variables)
        ret = self._eval_node(method_uast)
        self.variables.pop()
        return ret

    def _eval_unary_expr(self, node: UnaryExpr) -> ConcreteValueType:
        value = self._eval_node(node.expression)
        match node:
            case UnaryExpr(sign="++", op_before_expr=False):
                l_value = self.eval_l_value(node.expression)

                self.assign_value(l_value, value + 1)
                return value
            case UnaryExpr(sign="!", expression=expr):
                return not value
            case UnaryExpr(sign="*", op_before_expr=True):
                assert isinstance(value, PointerValue)
                return value.pointee.value
            case _:
                raise NotImplementedError(node)

    def _eval_binary_expr(self, node: BinaryExpr) -> ConcreteValueType:
        lv = self._eval_node(node.lhs)
        rv = self._eval_node(node.rhs)
        match node.operator:
            case "&&":
                return lv and rv
            case "||":
                return lv or rv
            case "^":
                assert isinstance(lv, int) and isinstance(rv, int)
                return lv ^ rv
            case ">":
                return lv > rv
            case "<":
                return lv < rv
            case "==":
                return lv == rv
            case "%":
                return lv % rv
            case "/":
                match (lv, rv):
                    case (int(), int()):
                        return int(lv / rv)
                    case _:
                        return lv / rv
            case "*":
                return lv * rv
            case "+":
                return lv + rv
            case "-":
                return lv - rv
            case _:
                raise NotImplementedError(stmt.operator)

    def _eval_compound_decl(self, node: CompoundDecl):
        for stmt in node.decls:
            self._eval_node(stmt)

    def _eval_var_decl(self, node: VarDecl):
        r_value = (
            self._eval_node(node.initializer)
            if node.initializer is not None
            else create_concreate_value_by_type(node.type)
        )
        var = self.variables[-1][node.variable.id] = Variable(None)
        if isinstance(r_value, ArrayInitializerValue):
            match node.type:
                case StructType():
                    r_value_new = StructValue()
                    for i, field in enumerate(node.type.fields):
                        r_value_new.fields[field.id] = r_value.values[i]()
                    raise NotImplementedError(node.type)
                case _:
                    raise NotImplementedError(node.type)
        self.assign_value(var, r_value)

    def _eval_method_decl(self, node: MethodDecl) -> List[ConcreteValueType]:
        for argument in node.type.pos_args:
            assert argument.name.id in self.variables[-1]
        cfg = CFGBuilder().build(node)
        current_block = cfg.entry_block
        while True:
            try:
                for stmt in current_block.statements:
                    self._eval_node(stmt)
                next_block = cfg.get_next_blocks(current_block)
                assert len(next_block) <= 1
                if len(next_block) == 0:
                    break
                current_block = next_block[0]
            except FunctionReturn as e:
                assert len(e.ret_value) == 1
                return e.ret_value[0]
        return None

    def _eval_assignment(self, node: Assignment):
        assert len(node.lhs) == 1
        assert len(node.rhs) == 1
        l_value = self.eval_l_value(node.lhs[0])
        r_value = self._eval_node(node.rhs[0])
        original_value = self._eval_node(node.lhs[0])
        match node.operator:
            case "=":
                pass
            case "+=":
                r_value = original_value + r_value
            case "-=":
                r_value = original_value - r_value
            case "*=":
                r_value = original_value * r_value
            case "/=":
                r_value = original_value / r_value
            case _:
                raise NotImplementedError(
                    f"Unsupported assignment operator: {node.operator}"
                )
        self.assign_value(l_value, r_value)
        return r_value

    def _eval_return_stmt(self, node: ReturnStmt):
        raise FunctionReturn([self._eval_node(item) for item in node.result])

    def _eval_node(self, node: SourceElement):
        if node.__class__ in self.evaluator_mapping:
            return self.evaluator_mapping[node.__class__](node)
        else:
            raise NotImplementedError(f"Node {node.__class__} is not implemented")

    def execute(
        self,
        scope_uast: MethodDecl,
    ) -> Variable:
        """
        Execute the UAST
        """
        return self._eval_node(scope_uast)
